import time
import functools
import psutil
from typing import (
    Callable,
    Any, Dict,
    Optional
)
from .logger import logger

class PerformanceMonitor:
    """Monitor performance metrics for ShellGPT."""
    
    def __init__(self):
        self.metrics: Dict[str, Dict[str, Any]] = {}
        self.start_time = time.time()
    
    def start_timer(self, name: str) -> None:
        """Start timing an operation."""
        self.metrics[name] = {
            'start_time': time.time(),
            'memory_start': psutil.Process().memory_info().rss
        }
        logger.debug(f"Started timing: {name}")
    
    def end_timer(self, name: str) -> Dict[str, Any]:
        """End timing an operation and return metrics."""
        if name not in self.metrics:
            logger.warning(f"No timer found for: {name}")
            return {}
        
        end_time = time.time()
        memory_end = psutil.Process().memory_info().rss
        
        metrics = self.metrics[name]
        duration = end_time - metrics['start_time']
        memory_used = memory_end - metrics['memory_start']
        
        result = {
            'duration': duration,
            'memory_used': memory_used,
            'memory_used_mb': memory_used / 1024 / 1024
        }
        
        logger.info(f"Operation {name} completed in {duration:.3f}s, "
                   f"memory used: {result['memory_used_mb']:.2f}MB")
        
        # Clean up
        del self.metrics[name]
        return result
    
    def get_system_info(self) -> Dict[str, Any]:
        """Get current system resource information."""
        process = psutil.Process()
        memory_info = process.memory_info()
        
        return {
            'cpu_percent': process.cpu_percent(),
            'memory_rss': memory_info.rss,
            'memory_rss_mb': memory_info.rss / 1024 / 1024,
            'memory_vms': memory_info.vms,
            'memory_vms_mb': memory_info.vms / 1024 / 1024,
            'num_threads': process.num_threads(),
            'create_time': process.create_time()
        }

# Global performance monitor
performance_monitor = PerformanceMonitor()

def monitor_performance(func_name: Optional[str] = None):
    """
    Decorator to monitor function performance.
    
    Args:
        func_name: Custom name for the function (optional)
    """
    def decorator(func: Callable) -> Callable:
        @functools.wraps(func)
        def wrapper(*args, **kwargs) -> Any:
            name = func_name or func.__name__
            performance_monitor.start_timer(name)
            
            try:
                result = func(*args, **kwargs)
                return result
            finally:
                performance_monitor.end_timer(name)
        
        return wrapper
    return decorator

def get_memory_usage() -> Dict[str, float]:
    """
    Get current memory usage information.
    
    Returns:
        Dict containing memory usage in MB
    """
    process = psutil.Process()
    memory_info = process.memory_info()
    
    return {
        'rss_mb': memory_info.rss / 1024 / 1024,
        'vms_mb': memory_info.vms / 1024 / 1024,
        'percent': process.memory_percent()
    }

def log_performance_stats() -> None:
    """Log current performance statistics."""
    system_info = performance_monitor.get_system_info()
    memory_usage = get_memory_usage()
    
    logger.info(f"Performance stats - "
               f"CPU: {system_info['cpu_percent']:.1f}%, "
               f"Memory: {memory_usage['rss_mb']:.2f}MB "
               f"({memory_usage['percent']:.1f}%)")

def check_system_resources() -> bool:
    """
    Check if system has sufficient resources.
    
    Returns:
        bool: True if resources are sufficient
    """
    memory_usage = get_memory_usage()
    
    # Check if memory usage is reasonable (less than 80% of available)
    if memory_usage['percent'] > 80:
        logger.warning(f"High memory usage: {memory_usage['percent']:.1f}%")
        return False
    
    return True
